import pySALEPlot as psp
from pylab import figure,arange,colorbar,loadtxt,figtext,shape
# This allows the figure to be subdivided nonuniformly
from matplotlib import gridspec
# Need this for the colorbars we will make on the mirrored plot
from mpl_toolkits.axes_grid1 import make_axes_locatable

import matplotlib.cm as cm
from matplotlib.colors import LinearSegmentedColormap

# This is how we create a nice greyscale colormap
cdict = cm.get_cmap('Spectral')._segmentdata
cdict['red'][0] = (0., 1., 1.)
cdict['blue'][0] = (0., 1., 1.)
cdict['green'][0] = (0., 1., 1,)
my_scmp = LinearSegmentedColormap('my_scmp', cdict)

cdict = cm.get_cmap('Greys')._segmentdata
gscale = [1.,0.9125,0.825,0.7375,0.645,0.5625,0.475,0.3875,0.3]
for i in range(9):
    scale = 0. + 0.125*i
    cdict['red'][i] = (scale, gscale[i], gscale[i])
    cdict['blue'][i] = cdict['red'][i] 
    cdict['green'][i] = cdict['red'][i]
    
my_cmap = LinearSegmentedColormap('my_cmap', cdict)

# Make an output directory
dirname='Plots'
psp.mkdir_p(dirname)

# Control parameters
simulations = ["Dilatancy"] # Specify the set of simulations to process (just one for the example)
isim = [100]        # timestep to make the plot from
xmax = [100]        # x-limit of region to plot
gravity_anomaly = 1 # To turn on a plot of the gravity anomaly
plot_cbar = 1       # To plot the colorbar
pmax = 17.          # Maximum porosity

# Specify the porosity contours to plot and how to label them
cfmt = {}
clevels = [0.1,1,5,10,15]
cstrs = ["0.1","1","5","10","15"]
for l,s in zip(clevels,cstrs):
   cfmt[l]=s

n=0
for sim in simulations:

    # Define datafilename
    datafilename = 'jdata.dat'

    # Open the datafile
    model=psp.opendatfile(datafilename)

    # Set the distance units to km
    model.setScale('km')

    # Set up a pylab figure
    if (gravity_anomaly == 1):
        # Specifiy the figure size (in inches)
	if (plot_cbar == 1):
            fig=figure(figsize=(8,10))
	else:
            fig=figure(figsize=(8,8))
        # Divide the figure into two rows, one column, with different heights
        gs = gridspec.GridSpec(2,1,height_ratios=[1,4])
        # Create axes with equal aspect in the second row
        ax=fig.add_subplot(gs[1],aspect='equal')
    else:
        # If a single image with no gravity profile, create a simple single axes figure.
	if (plot_cbar == 1):
            fig=figure(figsize=(8,7))
	else:
            fig=figure(figsize=(8,6))
        gs = gridspec.GridSpec(1,1)
        ax=fig.add_subplot(gs[0],aspect='equal')

    # Set the axis labels
    ax.set_xlabel('Distance from symmetry axis (km)')
    ax.set_ylabel('Height (km)')

    # Set the axis limits
    ax.set_xlim([0,100])
    ax.set_ylim([-40,10])

    # Read distension at the time step 'i' from the datafile
    i=isim[n]
    step=model.readStep('Alp',i)

    # Porosity in % from distension
    porosity=100.*(1.-1./step.data[0])

    # Plot the distension field as porosity as a colourmap and with contours
    print "Plotting porosity distribution for simulation: ",sim
    p=ax.pcolormesh(model.x,model.y,porosity,cmap=my_cmap,vmin=0.,vmax=pmax)
    q=ax.contour(model.xc,model.yc,porosity,linewidths=1,linestyles="solid",colors='k',levels=clevels,inline=1)
    ql=ax.clabel(q,inline=1,fontsize=12,fmt=cfmt)

    # And the material boundary

    b=ax.contour(model.xc,model.yc,step.cmc[0],colors='k',levels=[0.1])
    c=ax.contour(model.xc,model.yc,step.cmc[1],colors='k',levels=[0.1])
   
    # Add a colourbar, but only need to do this once
    if (plot_cbar == 1):
        cb=fig.colorbar(p,orientation='horizontal',shrink=0.7,pad=0.1)
        cb.set_label('Porosity (%)')

    # Add a tracer overlay
    for u in range(model.tracer_numu):
        tru=model.tru[u]
        # Plot the tracers in horizontal lines, every 10 lines
        for l in arange(0,len(tru.xlines),15):
            ax.plot(step.xmark[tru.xlines[l]],
                    step.ymark[tru.xlines[l]],
                    c='k',marker='.',linestyle='None',markersize=0.25)
        # Plot the tracers in vertical lines, every 10 lines
        for l in arange(0,len(tru.ylines),15):
            ax.plot(step.xmark[tru.ylines[l]],
                    step.ymark[tru.ylines[l]],
                    c='k',marker='.',linestyle='None',markersize=0.25)

    # Add a gravity anomaly plot if required
    if (gravity_anomaly == 1):
        ax2=fig.add_subplot(gs[0])
	# Calculate the gravity anomaly using pySALEPlot's inbuilt function
	# Note this returns a 4D array: [x,y] location of anomaly and [gx,gy] components of anomaly.
	grav_anomaly = model.gravityAnomaly(anomalyTime=i,anomaly="BOUGUER",method="Alp",ref_dist=[1,1.],ref_den=[2870.,3310.],bouguer_den=2870.,altitude=-1000., upperdepth=4000)
        # Plot the gravity profile gy as a function of x
	print "Plotting gravity anomaly for simulation: ",sim
        gp=ax2.plot(grav_anomaly[:,0],grav_anomaly[:,3],'k-')
        print grav_anomaly[:,0], grav_anomaly[:,3]
	# Other plotting stuff
        #ax2.set_ylim([-60,0])
        ax2.set_xlim([0,xmax[n]])
        ax2.set_ylabel('Bouguer anomaly (mGal)')
        ax2.grid(True)
        ax2.locator_params(axis = 'y',nbins=5)

    # Save the figure if one per simulation
    plotfilename = dirname+'/'+sim+'_final_porosity.png'
    print 'Saving plot as: ',plotfilename
    fig.savefig(plotfilename,dpi=300)

    n=n+1
